import pyvisa
import time

class N5171B:
    def __init__(self, ip_address):
        # 创建VISA资源管理器
        self.rm = pyvisa.ResourceManager()
        # 连接到信号源
        self.address = f'TCPIP0::{ip_address}::inst0::INSTR'
        try:
            self.instrument = self.rm.open_resource(self.address)
            # 设置超时时间（毫秒）
            self.instrument.timeout = 5000
            print("成功连接到信号源")
        except Exception as e:
            print(f"连接失败: {str(e)}")
            raise

    def get_amplitude(self):
        """获取输出幅值（单位：dBm）"""
        try:
            amplitude = self.instrument.query(':POW:LEV?')
            return float(amplitude)
        except Exception as e:
            print(f"获取幅值失败: {str(e)}")
            return None

    def get_frequency(self):
        """获取输出频率（单位：Hz）"""
        try:
            frequency = self.instrument.query(':FREQ:CW?')
            return float(frequency)
        except Exception as e:
            print(f"获取频率失败: {str(e)}")
            return None

    def set_amplitude(self, amplitude):
        """设置输出幅值（单位：dBm）
        参数:
            amplitude: 幅值，范围 -20 到 +30 dBm
        """
        try:
            self.instrument.write(f':POW:LEV {amplitude}')
            print(f"已设置幅值为: {amplitude} dBm")
        except Exception as e:
            print(f"设置幅值失败: {str(e)}")

    def set_frequency(self, frequency):
        """设置输出频率（单位：Hz）
        参数:
            frequency: 频率，范围 9kHz 到 6GHz
        """
        try:
            self.instrument.write(f':FREQ:CW {frequency}')
            print(f"已设置频率为: {frequency/1e6} MHz")
        except Exception as e:
            print(f"设置频率失败: {str(e)}")

    def close(self):
        """关闭连接"""
        if hasattr(self, 'instrument'):
            self.instrument.close()
        self.rm.close()

def main():
    # 信号源IP地址
    ip_address = '192.168.31.26'
    
    try:
        # 创建信号源实例
        signal_generator = N5171B(ip_address)
        
        # 首先读取当前设置
        amplitude = signal_generator.get_amplitude()
        frequency = signal_generator.get_frequency()
        print("初始设置值：")
        print(f"当前幅值: {amplitude:.2f} dBm")
        print(f"当前频率: {frequency/1e6:.2f} MHz")
        
        print("\n设置新的参数...")
        # 设置新的幅值和频率
        signal_generator.set_amplitude(0)  # 设置幅值为1 dBm
        signal_generator.set_frequency(81.25e6)  # 设置频率为81.25 MHz
        
        print("\n确认新的设置值：")
        # 再次读取确认设置
        amplitude = signal_generator.get_amplitude()
        frequency = signal_generator.get_frequency()
        print(f"当前幅值: {amplitude:.2f} dBm")
        print(f"当前频率: {frequency/1e6:.2f} MHz")
        
        # 关闭连接
        signal_generator.close()
        
    except Exception as e:
        print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main() 